import networkx as nx
import re
from typing import Dict, List, Any, Optional
from app.graph_rag.financial_kg import FinancialKnowledgeGraph
from app.agent.llm import get_llm
from ..core.logger import get_logger

logger = get_logger('graph_rag.query_engine')

class GraphRAGQuery:
    """GraphRAG查询器"""
    
    def __init__(self):
        """
        初始化GraphRAG查询器
        
        Args:
            knowledge_graph: 金融知识图谱实例
        """
        self.knowledge_graph = knowledge_graph
        self.llm = get_llm()
    
    def query_financial_relationships(self, query: str) -> str:
        """
        查询金融实体关系
        
        Args:
            query: 用户查询
            
        Returns:
            格式化的回答
        """
        try:
            # 识别查询中的实体
            entities = self._extract_entities(query)
            
            if not entities:
                # 如果无法提取实体，尝试使用更通用的方法
                return self._handle_generic_query(query)
            
            # 获取相关子图
            subgraph = self.knowledge_graph.get_subgraph(entities, depth=2)
            
            if subgraph.number_of_nodes() == 0:
                return f"未找到与 {', '.join(entities)} 相关的信息。"
            
            # 格式化子图信息
            graph_info = self._format_subgraph_info(subgraph)
            
            # 构建提示词
            prompt = self._build_prompt(query, graph_info)
            
            # 使用LLM生成回答
            response = self.llm.invoke(input=prompt)
            return response
            
        except Exception as e:
            logger.error(f"GraphRAG查询时出错: {e}")
            return f"查询过程中出现错误: {str(e)}"
    
    def _handle_generic_query(self, query: str) -> str:
        """
        处理通用查询
        
        Args:
            query: 用户查询
            
        Returns:
            格式化的回答
        """
        try:
            # 构建通用查询提示词
            prompt = f"""
你是一个金融数据分析专家，能够回答各种金融相关问题。

用户问题：
{query}

请基于你的金融知识，用中文简洁准确地回答用户问题。如果问题涉及具体数据但没有提供足够的上下文信息，请说明需要哪些额外信息才能准确回答。
回答时请注意：
1. 使用清晰的结构化格式
2. 重点突出关键信息
3. 避免编造未提供的信息
4. 回答应直接针对用户问题，不要添加无关信息

回答：
"""
            
            # 使用LLM生成回答
            response = self.llm.invoke(input=prompt)
            return response
            
        except Exception as e:
            logger.error(f"处理通用查询时出错: {e}")
            return f"处理通用查询时出现错误: {str(e)}"
    
    def _extract_entities(self, query: str) -> List[str]:
        """
        从查询中提取金融实体
        
        Args:
            query: 用户查询
            
        Returns:
            实体列表
        """
        # 简单的实体提取方法，实际应用中可以使用NER模型
        entities = {
            "fund_codes": [],
            "stock_codes": [],
            "company_names": []  # 新增公司名称实体列表
        }
        
        # 提取基金代码
        fund_pattern = r'[fF][0-9]{6}'
        entities["fund_codes"] = re.findall(fund_pattern, query)
        
        # 提取股票代码
        stock_pattern = r'[sS][hzHZ][0-9]{6}'
        entities["stock_codes"] = re.findall(stock_pattern, query)
        # 新增：提取公司名称实体
        company_patterns = [
            r'([^，。,;；\s]+?[股份有限]{2,4}公司)',
            r'([^，。,;；\s]+?有限责任公司)'
        ]
        
        # 定义股票模式列表
        stock_patterns = [
            r'([0-9]{6}\.[Ss][Hh])',  # 上交所股票代码格式，如 600000.SH
            r'([0-9]{6}\.[Ss][Zz])'   # 深交所股票代码格式，如 000001.SZ
        ]

        for pattern in stock_patterns:
            matches = re.findall(pattern, query)
            for match in matches:
                # 简化处理，实际应根据后缀判断市场
                code = match.split('.')[0]
                entities["stock_codes"].append(f"stock_{code}")
        for pattern in company_patterns:
            matches = re.findall(pattern, query)
            if matches:
                entities["company_names"].extend(matches)

        # 将字典中的所有实体合并为一个列表返回
        all_entities = []
        for entity_list in entities.values():
            all_entities.extend(entity_list)
            
        return all_entities
    
    def _format_subgraph_info(self, subgraph: nx.DiGraph) -> str:
        """
        格式化子图信息
        
        Args:
            subgraph: 子图
            
        Returns:
            格式化后的信息字符串
        """
        if subgraph.number_of_nodes() == 0:
            return "无相关信息"
        
        info_lines = []
        info_lines.append(f"找到 {subgraph.number_of_nodes()} 个实体和 {subgraph.number_of_edges()} 个关系：")
        
        # 添加节点信息
        info_lines.append("\n实体信息：")
        for node, data in subgraph.nodes(data=True):
            node_type = data.get('type', 'unknown')
            name = data.get('name', 'unknown')
            # 添加更多节点属性信息
            attrs = {k: v for k, v in data.items() if k not in ['type', 'name']}
            info_lines.append(f"  - {node_type}: {name} ({node})")
            if attrs:
                attr_str = ", ".join([f"{k}: {v}" for k, v in attrs.items()])
                info_lines.append(f"    属性: {attr_str}")
        
        # 添加边信息
        info_lines.append("\n关系信息：")
        for source, target, data in subgraph.edges(data=True):
            relation = data.get('relation', 'related_to')
            source_name = subgraph.nodes[source].get('name', source)
            target_name = subgraph.nodes[target].get('name', target)
            info_lines.append(f"  - {source_name} --[{relation}]--> {target_name}")
            
            # 添加边的属性信息
            attrs = {k: v for k, v in data.items() if k != 'relation'}
            if attrs:
                attr_str = ", ".join([f"{k}: {v}" for k, v in attrs.items()])
                info_lines.append(f"    属性: {attr_str}")
        
        return "\n".join(info_lines)
    
    def _build_prompt(self, query: str, graph_info: str) -> str:
        """
        构建提示词
        
        Args:
            query: 用户查询
            graph_info: 图谱信息
            
        Returns:
            构建的提示词
        """
        prompt = f"""
你是一个金融数据分析专家，能够理解和解释金融实体之间的复杂关系。

基于以下知识图谱信息回答用户的问题：

知识图谱信息：
{graph_info}

用户问题：
{query}

请根据提供的知识图谱信息，用中文简洁准确地回答用户问题。如果信息不足以回答问题，请说明原因。
回答时请注意：
1. 使用清晰的结构化格式
2. 重点突出关键关系和数据
3. 避免编造图谱中未提供的信息
4. 如果有多个相关实体，请分别说明
5. 回答应直接针对用户问题，不要添加无关信息

回答：
"""
        return prompt.strip()
    
    def find_investment_patterns(self, fund_code: str) -> str:
        """
        分析基金的投资模式
        
        Args:
            fund_code: 基金代码
            
        Returns:
            投资模式分析
        """
        try:
            fund_node = f"fund_{fund_code}"
            if fund_node not in self.knowledge_graph.graph.nodes:
                return f"未找到基金 {fund_code} 的信息。"
            
            # 获取基金基本信息
            fund_info = dict(self.knowledge_graph.graph.nodes[fund_node])
            
            # 获取基金持仓信息
            holdings = []
            for neighbor in self.knowledge_graph.graph.neighbors(fund_node):
                edge_data = self.knowledge_graph.graph.get_edge_data(fund_node, neighbor)
                neighbor_info = dict(self.knowledge_graph.graph.nodes[neighbor])
                
                if edge_data.get('relation') == 'holds':
                    holdings.append({
                        'stock': neighbor_info,
                        'market_value': edge_data.get('market_value'),
                        'ratio': edge_data.get('ratio')
                    })
            
            # 按持仓比例排序
            holdings.sort(key=lambda x: float(x['ratio']) if x['ratio'] else 0, reverse=True)
            
            # 构建分析提示词
            holdings_info = "\n".join([
                f"  - {h['stock'].get('name', 'unknown')} (比例: {h['ratio']}, 市值: {h['market_value']})"
                for h in holdings[:10]  # 限制显示前10个持仓
            ])
            
            prompt = f"""
分析以下基金的投资模式：

基金信息：
- 名称: {fund_info.get('name', 'unknown')}
- 类型: {fund_info.get('fund_type', 'unknown')}
- 管理人: {fund_info.get('manager', 'unknown')}

主要持仓：
{holdings_info}

请基于以上信息分析该基金的投资模式，包括：
1. 投资偏好（如行业偏好、市值偏好等）
2. 风险特征
3. 投资策略特点

分析：
"""
            
            response = self.llm.invoke(input=prompt)
            return response
            
        except Exception as e:
            logger.error(f"分析投资模式时出错: {e}")
            return f"分析过程中出现错误: {str(e)}"